import os
import cv2
import torch
from movenet.opts import opts
from movenet.detectors.detector_factory import detector_factory

image_ext = ["jpg", "jpeg", "png", "webp"]
video_ext = ["mp4", "mov", "avi", "mkv"]
time_stats = ["tot", "load", "pre", "net", "dec", "post", "merge"]


def demo(opt):
    os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus_str
    opt.debug = max(opt.debug, 1)
    Detector = detector_factory[opt.task]
    detector = Detector(opt)

    if opt.demo == "webcam" or opt.demo[opt.demo.rfind(".") + 1 :].lower() in video_ext:
        cam = cv2.VideoCapture(0 if opt.demo == "webcam" else opt.demo)
        detector.pause = False
        while True:
            success, img = cam.read()
            if not success:
                print("no more images, close.")
                return
            if opt.debug < 4:
                cv2.imshow("input", img)
            ret = detector.run(img)
            time_str = ""
            for stat in time_stats:
                time_str = time_str + "{} {:.3f}s |".format(stat, ret[stat])
            print(time_str)
            if opt.debug < 4 and cv2.waitKey(1) == 27:
                return  # esc to quit
    else:
        if os.path.isdir(opt.demo):
            image_names = []
            ls = os.listdir(opt.demo)
            for file_name in sorted(ls):
                ext = file_name[file_name.rfind(".") + 1 :].lower()
                if ext in image_ext:
                    image_names.append(os.path.join(opt.demo, file_name))
        else:
            image_names = [opt.demo]

        for image_name in image_names:
            ret = detector.run(image_name)
            time_str = ""
            for stat in time_stats:
                time_str = time_str + "{} {:.3f}s |".format(stat, ret[stat])
            print(time_str)


if __name__ == "__main__":
    opt = opts().init()
    demo(opt)
